package main.java.com.itlin.datastructure.redblacktree;

public class RedBlackTree {
    public boolean contains(int key) {

        Node node = root;
        while (node != null) {
            if (key == node.key) {
                return true;
            } else if (key < node.key) {
                node = node.left;
            } else {
                node = node.right;
            }
        }
        return false;
    }



    static enum Color {
        RED, BLACK;
    }

    public static class Node{
         public Object value;
         int key;
         Node left;
         Node right;
         Node parent; // 父节点
         Color color = Color.BLACK; // 颜色

        public Node(int key, Object value) {
            this.key = key;
            this.value = value;
        }

        // 工具方法

        // 是否是左孩子
        boolean isLeftChild(){
            return parent != null && this == parent.left;
        }
        // 是否是右孩子
        public boolean isRightChild() {
            return parent != null && this == parent.right;
        }


        // 找当前节点的叔叔
        Node uncle(){
            if(parent == null || parent.parent == null){
                return null;
            }
            // 父节点是左孩子
            if(parent.isLeftChild()){
                return parent.parent.right;
            }else{
                return parent.parent.left;
            }
        }

        // 兄弟
        Node sibling(){
            if(parent == null){
                return null;
            }
            // 是父节点左孩子
            if(isLeftChild()){
                return parent.right;
            }else{
                return parent.left;
            }
        }

    }

    private Node root;

    //判断节点是红色还是黑色
    boolean isRed(Node node){
        return node != null && node.color == Color.RED;
    }

    boolean isBlack(Node node){
        return node == null ||node.color == Color.BLACK;
    }

    // 旋转  node 是要旋转的节点
    //右旋 1.parent 处理  2. 旋转后的新根 在方法内直接建立好
    private void rightRotate(Node node){
        Node parent = node.parent;
        Node child = node.left; //旋转后的新根节点（针对旋转部份的根节点）
        Node childSon = child.right;// 要换父节点的节点
        if (childSon != null){
            childSon.parent = node;
        }
        child.right = node;
        child.parent = parent;
        node.left = childSon;
        node.parent = child;
        if(parent == null){
            root = child;
        }else if(parent.left == node){
            parent.left = child;
        }else{
            parent.right = child;
        }
    }

    private void leftRotate(Node node){
        Node parent = node.parent;
        Node child = node.right;
        Node childSon = child.left;
        if(childSon != null){
            childSon.parent = node;
        }
        child.left = node;
        child.parent = parent;
        node.right = childSon;
        node.parent = child;
        if(parent == null){
            root = child;
        }else if(parent.left == node){
            parent.left = child;
        }else{
            parent.right = child;
        }
    }

    /**
     * 新增或更新
     * 正常增、遇到红红不平衡的情况进行调整
     * @param key - 键
     * @param value - 值
     */
    public void put(int key, Object value){
        //先找到节点的位置
        Node p = root;
        Node parent = null;//记录要新增位置的父节点
        while(p != null){
            parent = p;
            if(key < p.key){
                p = p.left;
            }else if(key > p.key){
                p = p.right;
            }else {
                p.value = value;//更新
                return;
            }
        }
        // 新增
        Node inserted = new Node(key, value);
        if(parent == null){
            root = inserted;
        }
        else if (key < parent.key){
            parent.left = inserted;
            inserted.parent = parent;
        }
        else {
            parent.right = inserted;
            inserted.parent = parent;
        }
        fixRedRed(inserted);

    }

    void fixRedRed(Node x){
        // case 1 插入节点是根节点  将颜色变黑
        if(x == root){
            x.color = Color.BLACK;
            return;
        }
        // case 2 插入节点父亲是黑色， 无需调整
        if(isBlack(x.parent)){
            return;
        }
        // case 3 插入节点父亲是红色，叔叔节点是红色
        Node parent = x.parent;
        Node uncle = x.uncle();
        Node grandpa = parent.parent;
        if(isRed(uncle)){
            // 叔叔节点变黑
            uncle.color = Color.BLACK;
            parent.color = Color.BLACK;
            grandpa.color = Color.RED;
            fixRedRed(grandpa); // 递归调用
            return;
        }
        // case 4 插入节点父亲是红色，叔叔节点是黑色
        // 插入节点是父亲(本身也是祖父的左孩子)的左孩子
        if(parent.isLeftChild() && x.isLeftChild()){
            parent.color = Color.BLACK;
            grandpa.color = Color.RED;
            rightRotate(grandpa);
        }
        // 插入节点是父亲(本身也是祖父的左孩子)的右孩子
        else if (parent.isLeftChild() && x.isRightChild()){
            // 在 父亲处左旋
            leftRotate(parent);
            x.color = Color.BLACK;
            parent.color = Color.RED;
            rightRotate(grandpa);
        }
        // 插入节点是父亲(本身也是祖父的右孩子)的左孩子
        else if (parent.isRightChild() && x.isLeftChild()){
            // 在 父亲处右旋
            rightRotate(parent);
            x.color = Color.BLACK;
            parent.color = Color.RED;
            leftRotate(grandpa);
        }
        // 插入节点是父亲(本身也是祖父的右孩子)的右孩子
        else if (parent.isRightChild() && x.isRightChild()){
            parent.color = Color.BLACK;
            grandpa.color = Color.RED;
            leftRotate(grandpa);
        }
    }

    /**
     * 删除
     * 正常删除、会用到李代桃僵技巧、遇到红红不平衡的情况进行调整
     * @param key - 键
     */
    public void remove(int key){
        Node deleted = findKey(key);
        if(deleted == null){
            return;
        }
        doRemove(deleted);

    }

    //
    private void fixDoubleBlack(Node x){
        if(x == root){
            return;
        }
        Node parent = x.parent;
        Node sibling = x.sibling();
        // case 3 兄弟节点是红色
        if(isRed(sibling)){
            if (x.isLeftChild()){
                leftRotate(parent);
            } else {
                rightRotate(parent);

            }
            parent.color = Color.RED;
            sibling.color = Color.BLACK;
            fixDoubleBlack(x);// 递归调用case 4
            return;
        }
        // case 4 兄弟节点是黑色
        if (sibling != null){
            if(isBlack(sibling.left) && isBlack(sibling.right)){
                // 兄弟节点是黑色，且兄弟节点的左右孩子都是黑色

                sibling.color = Color.RED;
                if(isRed(parent)){
                    parent.color = Color.BLACK;
                }
                fixDoubleBlack(parent);

            }// case 5 兄弟节点是黑色， 侄子有红色
            else {
                // 兄弟节点 是左孩子， 左侄子是红色  LL不平衡
                if(sibling.isLeftChild() && isRed(sibling.left)){
                    rightRotate(sibling);
                    sibling.left.color = Color.BLACK;
                    sibling.color = parent.color;
                }
                // 兄弟节点 是右孩子， 右侄子是红色  RR不平衡
                else if(sibling.isRightChild() && isRed(sibling.right)){
                    leftRotate(sibling);
                    sibling.right.color = Color.BLACK;
                    sibling.color = parent.color;
                }
                // 兄弟节点 是左孩子， 右侄子是红色  LR不平衡
                else if (sibling.isLeftChild() && isRed(sibling.right)){
                    sibling.right.color = parent.color;
                    leftRotate(sibling);
                    rightRotate(parent);
                }
                // 兄弟节点 是右孩子， 左侄子是红色  RL不平衡
                else if (sibling.isRightChild() && isRed(sibling.left)){
                    sibling.left.color = parent.color;
                    rightRotate(sibling);
                    leftRotate(parent);
                }
                parent.color = Color.BLACK;
            }
        }// 兄弟节点是null 为黑色
        else {
            fixDoubleBlack(parent);
        }

    }

    private void doRemove(Node deleted) {
        Node replaced = findReplaced(deleted);//找到剩余节点
        Node parent = deleted.parent;
        if(replaced == null){
            //删除节点没有孩子
            // case1 删除节点是根节点
            if(deleted == root){
                root = null;
            }// case2 删除节点不是根节点
            else {
                if(isBlack(deleted)){
                    // 复杂调整
                    fixDoubleBlack(deleted);
                } else {
                    // 红色叶子，不用处理
                }
                if (deleted.isLeftChild()){
                    parent.left = null;
                }else {
                    parent.right = null;
                }
                deleted.parent = null;
            }
            return ;
        }
        // 删除节点有一个孩子
        if(deleted.left == null || deleted.right == null){
            // 删除节点是根节点
            if(root == deleted){
                root.key = replaced.key;
                root.value = replaced.value;
                root.left = root.right = null;
            }
            // 删除节点不是根节点
            else {
                if(deleted.isLeftChild()){
                    parent.left = replaced;
                }else {
                    parent.right = replaced;
                }
                replaced.parent = parent;
                deleted.right = deleted.left = deleted.parent = null;
                if(isBlack(deleted) && isBlack(replaced)){
                    // 复杂调整
                    fixDoubleBlack(replaced);
                }else {
                    // 红色叶子，不用处理 case 2
                    replaced.color = Color.BLACK;
                }
            }
            return;
        }
        // 删除节点有两个孩子 == > 转换成只有一个孩子或没有孩子的情况
        // 李代桃僵   交换 删除节点和交换节点的  键和值
        int t = deleted.key;
        deleted.key = replaced.key;
        replaced.key = t;

        Object o = deleted.value;
        deleted.value = replaced.value;
        replaced.value = o;
        doRemove(replaced);
    }

    //查找节点
    public Node findKey(int key){
        Node p = root;
        while(p != null){
            if(key < p.key){
                p = p.left;
            }else if(key > p.key){
                p = p.right;
            }else {
                return p;
            }

        }
        return null;
    }

    //查找剩余节点
    Node findReplaced(Node deleted){
        if (deleted.right == null && deleted.left == null){
            return null;
        }
        if (deleted.right != null){
            return deleted.right;
        }
        if (deleted.left != null){
            return deleted.left;
        }
        // 左右节点都有的时候， 找后继节点
        Node s = deleted.right;
        while(s.left != null){
            s = s.left;
        }
        return s;
    }




}
